import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Scanner;

public class AES {
	public static final String ZEROS = "00000000";
	public static final String XOR_STR = "00011011";	// 84310
	
	public static final String[][] FIX_COLUMN = new String[][] {	// 列混淆变换
		{"02", "03", "01", "01"}, 
		{"01", "02", "03", "01"}, 
		{"01", "01", "02", "03"}, 
		{"03", "01", "01", "02"} 
	};
	
	public static final String[][] FIX_COLUMN_INVERSE = new String[][] {	// 列混淆逆向变换
		{"0E", "0B", "0D", "09"}, 
		{"09", "0E", "0B", "0D"}, 
		{"0D", "09", "0E", "0B"}, 
		{"0B", "0D", "09", "0E"} 
	};
	
	public static final String[] RC = new String[]{"01", "02", "04", "08", "10", "20", "40", "80", "1B", "36"};
	
	public static final String[][] S = new String[][]{
		{"63", "7c", "77", "7b", "f2", "6b", "6f", "c5", "30", "01", "67", "2b", "fe", "d7", "ab", "76"}, 
		{"ca", "82", "c9", "7d", "fa", "59", "47", "f0", "ad", "d4", "a2", "af", "9c", "a4", "72", "c0"},
		{"b7", "fd", "93", "26", "36", "3f", "f7", "cc", "34", "a5", "e5", "f1", "71", "d8", "31", "15"}, 
		{"04", "c7", "23", "c3", "18", "96", "05", "9a", "07", "12", "80", "e2", "eb", "27", "b2", "75"}, 
		{"09", "83", "2c", "1a", "1b", "6e", "5a", "a0", "52", "3b", "d6", "b3", "29", "e3", "2f", "84"}, 
		{"53", "d1", "00", "ed", "20", "fc", "b1", "5b", "6a", "cb", "be", "39", "4a", "4c", "58", "cf"}, 
		{"d0", "ef", "aa", "fb", "43", "4d", "33", "85", "45", "f9", "02", "7f", "50", "3c", "9f", "a8"},
		{"51", "a3", "40", "8f", "92", "9d", "38", "f5", "bc", "b6", "da", "21", "10", "ff", "f3", "d2"}, 
		{"cd", "0c", "13", "ec", "5f", "97", "44", "17", "c4", "a7", "7e", "3d", "64", "5d", "19", "73"}, 
		{"60", "81", "4f", "dc", "22", "2a", "90", "88", "46", "ee", "b8", "14", "de", "5e", "0b", "db"},
		{"e0", "32", "3a", "0a", "49", "06", "24", "5c", "c2", "d3", "ac", "62", "91", "95", "e4", "79"},
		{"e7", "c8", "37", "6d", "8d", "d5", "4e", "a9", "6c", "56", "f4", "ea", "65", "7a", "ae", "08"},
		{"ba", "78", "25", "2e", "1c", "a6", "b4", "c6", "e8", "dd", "74", "1f", "4b", "bd", "8b", "8a"},
		{"70", "3e", "b5", "66", "48", "03", "f6", "0e", "61", "35", "57", "b9", "86", "c1", "1d", "9e"},
		{"e1", "f8", "98", "11", "69", "d9", "8e", "94", "9b", "1e", "87", "e9", "ce", "55", "28", "df"},
		{"8c", "a1", "89", "0d", "bf", "e6", "42", "68", "41", "99", "2d", "0f", "b0", "54", "bb", "16"}

	};
	
	public static final String[][] S_INVERSE = new String[][]{
		{"52", "09", "6a", "d5", "30", "36", "a5", "38", "bf", "40", "a3", "9e", "81", "f3", "d7", "fb"},
		{"7c", "e3", "39", "82", "9b", "2f", "ff", "87", "34", "8e", "43", "44", "c4", "de", "e9", "cb"},
		{"54", "7b", "94", "32", "a6", "c2", "23", "3d", "ee", "4c", "95", "0b", "42", "fa", "c3", "4e"},
		{"08", "2e", "a1", "66", "28", "d9", "24", "b2", "76", "5b", "a2", "49", "6d", "8b", "d1", "25"},
		{"72", "f8", "f6", "64", "86", "68", "98", "16", "d4", "a4", "5c", "cc", "5d", "65", "b6", "92"},
		{"6c", "70", "48", "50", "fd", "ed", "b9", "da", "5e", "15", "46", "57", "a7", "8d", "9d", "84"},
		{"90", "d8", "ab", "00", "8c", "bc", "d3", "0a", "f7", "e4", "58", "05", "b8", "b3", "45", "06"},
		{"d0", "2c", "1e", "8f", "ca", "3f", "0f", "02", "c1", "af", "bd", "03", "01", "13", "8a", "6b"},
		{"3a", "91", "11", "41", "4f", "67", "dc", "ea", "97", "f2", "cf", "ce", "f0", "b4", "e6", "73"},
		{"96", "ac", "74", "22", "e7", "ad", "35", "85", "e2", "f9", "37", "e8", "1c", "75", "df", "6e"},
		{"47", "f1", "1a", "71", "1d", "29", "c5", "89", "6f", "b7", "62", "0e", "aa", "18", "be", "1b"},
		{"fc", "56", "3e", "4b", "c6", "d2", "79", "20", "9a", "db", "c0", "fe", "78", "cd", "5a", "f4"},
		{"1f", "dd", "a8", "33", "88", "07", "c7", "31", "b1", "12", "10", "59", "27", "80", "ec", "5f"},
		{"60", "51", "7f", "a9", "19", "b5", "4a", "0d", "2d", "e5", "7a", "9f", "93", "c9", "9c", "ef"},
		{"a0", "e0", "3b", "4d", "ae", "2a", "f5", "b0", "c8", "eb", "bb", "3c", "83", "53", "99", "61"},
		{"17", "2b", "04", "7e", "ba", "77", "d6", "26", "e1", "69", "14", "63", "55", "21", "0c", "7d"}
		
	};
	
	
	public static void main(String[] args) {
		Scanner scanner = new Scanner(System.in);
		String choose = null;
		while (true) {
			System.out.println("AES算法, 请选择您要进行的操作(按下对应操作数字后回车): (1)加密 (2)解密 (3)退出");
			choose = scanner.nextLine();
			if ("1".equals(choose)) {
				System.out.print("请输入明文(不含中文): ");
				String plainText = scanner.nextLine();
				System.out.print("请输入密钥(16个字符以内, 不含中文): ");
				String secret = scanner.nextLine();
				String cipher = encrypt(plainText, secret);
				System.out.println("加密结果: " + cipher);
			} else if ("2".equals(choose)) {
				System.out.print("请输入密文(不含中文): ");
				String cipher = scanner.nextLine();
				System.out.print("请输入密钥(16个字符以内, 不含中文): ");
				String secret = scanner.nextLine();
				String plainText = decrypt(cipher, secret);
				System.out.println("解密结果: " + plainText);
			} else if ("3".equals(choose)) {
				System.exit(0);
			}
		}
		
	}
	
	public static String encrypt(String plainText, String secret) {
		// 0. 明文/密钥分组
		String[][][] plainMatrix = divide4PlainText(plainText);
		String[][] initSecret = divide4Secret(secret);
		StringBuilder sb = new StringBuilder(plainMatrix.length << 4); // 用于拼接密文
		// 1. 构造密钥组
		String[][][] extendSecret = extendSecret(initSecret);
		for (int i = 0; i < plainMatrix.length; i++) {
			// 2. 轮密钥加
			plainMatrix[i] = xor4HexArr(plainMatrix[i], extendSecret[0]);
			// 3.10轮重复加密操作
			for (int round = 1; round <= 10; round++) {
				// 字节代替
				replaceByte(plainMatrix[i], S);
				// 行移位
				plainMatrix[i] = moveRow(plainMatrix[i]);
				// 列混淆(第十轮不进行)
				if (round != 10) {
					plainMatrix[i] = fixColumn(plainMatrix[i], FIX_COLUMN);
				}
				// 轮密钥加
				plainMatrix[i] = xor4HexArr(plainMatrix[i], extendSecret[round]);
			}
			// 4.拼接密文
			for (int k = 0; k < 4; k++) {
				for (int j = 0; j < 4; j++) {
					sb.append(plainMatrix[i][j][k]);
				}
			}
		}
		return sb.toString();
	}

	/**
	 * 解密
	 * @param cipher 密文16进制字符串
	 * @param secret 密钥字符串
	 */
	public static String decrypt(String cipher, String secret) {
		// 0. 密文/密钥分组
		String[][][] cipherMatrix = divide4Hex(cipher);
		String[][] initSecret = divide4Secret(secret);
		List<Byte> list = new ArrayList<>(cipherMatrix.length << 4 + 16);
		// 1. 构造密钥组
		String[][][] extendSecret = extendSecret(initSecret);
		for (int i = 0; i < cipherMatrix.length; i++) {
			// 2. 轮密钥加
			cipherMatrix[i] = xor4HexArr(cipherMatrix[i], extendSecret[10]);
			// 3.10轮重复加密操作
			for (int round = 9; round >= 0; round--) {
				// 逆行移位
				cipherMatrix[i] = moveRowInverse(cipherMatrix[i]);
				// 逆字节代替
				replaceByte(cipherMatrix[i], S_INVERSE);
				// 轮密钥加
				cipherMatrix[i] = xor4HexArr(cipherMatrix[i], extendSecret[round]);
				// 逆列混淆(第十轮不进行)
				if (round != 0) {
					cipherMatrix[i] = fixColumn(cipherMatrix[i], FIX_COLUMN_INVERSE);
				}
			}
			// 4.还原明文
			for (int k = 0; k < 4; k++) {
				for (int j = 0; j < 4; j++) {
					byte byteVal = Byte.valueOf(cipherMatrix[i][j][k], 16);
					if (byteVal != 0) {
						list.add(byteVal);
					}
				}
			}
		}
		byte[] bytes = new byte[list.size()];
		for (int i = 0; i < bytes.length; i++) {
			bytes[i] = list.get(i);
		}
		return new String(bytes);
	}
	
	
	/**
	 * 打印4 * 4矩阵
	 * @param arr
	 */
	public static void print(String[][] arr) {
		for (int row = 0; row < 4; row++) {
			for (int col = 0; col < 4; col++) {
				System.out.print(arr[row][col] + " ");
			}
			System.out.println();
		}
		System.out.println("==============");
	}
	
	/**
	 * 明文
	 * 1. 将明文的普通字符串转化成对应的多个字节数组
	 * 2. 每个字节数组都不能超过16字节
	 * 3. 将字节数组按列排成4 * 4矩阵
	 * 4. 字节不足时用00充当
	 */
	public static String[][][] divide4PlainText(String plainText) {
		byte[] bytes = plainText.getBytes();
		int len = (bytes.length & 15) == 0  ? (bytes.length >>> 4) : (bytes.length >>> 4) + 1;
		String[][][] ret = new String[len][4][4];
		for (int k = 0; k < len; k++) {
			for (int i = 0; i < 4; i++) {
				for (int j = 0; j < 4; j++) {
					int index = (i << 2) + j + (k << 4);
					if (index < bytes.length) {
						String hex = Integer.toHexString(bytes[index]);
						if (hex.length() == 1) {
							hex = "0" + hex;
						}
						ret[k][j][i] = hex;
					} else {
						ret[k][j][i] = "00";
					}
				}
			}
		}
		return ret;
	}
	/**
	 * 密钥分组
	 * 1. 密钥的普通字符串转化成对应的字节数组
	 * 2. 不能超过16字节
	 * 3. 将字节数组按列排成4 * 4矩阵
	 * 4. 字节不足时用00充当
	 */
	public static String[][] divide4Secret(String secret) {
		String[][] ret = new String[4][4];
		byte[] bytes = secret.getBytes();
		if (bytes.length > 16) {
			throw new RuntimeException("密钥不能超过16个字节");
		}
		for (int i = 0; i < 4; i++) {
			for (int j = 0; j < 4; j++) {
				int index = (i << 2) + j;
				if (index < bytes.length) {
					String hex = Integer.toHexString(bytes[index]);
					if (hex.length() == 1) {
						hex = "0" + hex;
					}
					ret[j][i] = hex;
				} else {
					ret[j][i] = "00";
					
				}
			}
		}
		return ret;
	}
	

	/*
	 * 
	 * 将十六进制字符串分成多组4 * 4
	 * 如: 0123456789abcdeffedcba9876543210  (16字节)
	 * 转成: 
	 		01 89 fe 76 
			23 ab dc 54 
			45 cd ba 32 
			67 ef 98 10
	 * 转入的是密文(长度一定是32的倍数)
	 */
	public static String[][][] divide4Hex(String hexString) {
		int len = hexString.length() >>> 5;
		String[][][] ret = new String[len][4][4];
		int index = 0;
		for (int k = 0; k < len; k++) {
			for (int i = 0; i < 4; i++) {
				for (int j = 0; j < 4; j++) {
					ret[k][j][i] = hexString.substring(index, index + 2);
					index += 2;
				}
			}
		}
		return ret;
	}
	
	/**
	 * 密钥扩展
	 */
	public static String[][][] extendSecret(String[][] initSecret) {
		
		String[][] w = new String[44][4];
		// 1. w0-w3等于初始密钥
		for (int i = 0; i < initSecret.length; i++) {
			for (int j = 0; j < 4; j++) {
				w[i][j] = initSecret[j][i];
			}
		}
		// 2.计算w4-w43
		for (int i = 4; i < 44; i++) {
			String[] temp = w[i - 1];
			if ((i & 3) == 0) { // 如果wi下标i可以被4整除, 则w[i] = w[i - 4] xor z[i >>> 2]
				temp =  g(w[i - 1], (i >>> 2) - 1);
			} // 否则, w[i] = w[i - 1] xor w[i - 4]
			for (int j = 0; j < 4; j++) {
				w[i][j] = xor4Hex(w[i - 4][j], temp[j]);
			}
		}
		
		// 3.计算构造密钥组返回
		String[][][] secretGroup = new String[11][4][4];
		for (int i = 0; i < 11; i++) {
			for (int j = 0; j < 4; j++) {
				for (int k = 0; k < 4; k++) {
					secretGroup[i][j][k] = w[(i << 2) + k][j];
				}
			}
		}
		return secretGroup;
	}
	
	
	/**
	 * 辅助函数g
	 * @param w 当前轮最后一列
	 * @param round 当前轮数
	 * @return
	 */
	public static String[] g(String[] w, int round) {
		String[] z = new String[4];
		for (int i = 0; i < 4; i++) {
			// 1.字循环
			// 2.字代替
			z[i] = searchS(w[i + 1 & 3], S);
			// 3.xor轮常数
			if (i == 0) {
				z[i] = xor4Hex(z[i], RC[round]);
			}
		}
		return z;
	}
	
	/**
	 * 字节替换
	 */
	public static void replaceByte(String[][] arr, String[][] s) {
		for (int i = 0; i < 4; i++) {
			for (int j = 0; j < 4; j++) {
				arr[i][j] = searchS(arr[i][j], s);
			}
		}
	}
	
	/**
	 * 行移位
	 */
	public static String[][] moveRow(String[][] arr) {
		String[][] newArr = new String[4][4];
		for (int row = 0; row < 4; row++) {
			for (int col = 0; col < 4; col++) {
				newArr[row][col] = arr[row][(col + row) & 3];
			}
		}
		return newArr;
	}
	
	/**
	 * 逆行移位
	 */
	public static String[][] moveRowInverse(String[][] arr) {
		String[][] newArr = new String[4][4];
		for (int row = 0; row < 4; row++) {
			for (int col = 0; col < 4; col++) {
				newArr[row][col] = arr[row][(4 - row + col) & 3];
			}
		}
		return newArr;
	}
	
	/**
	 * 列混淆
	 */
	public static String[][] fixColumn(String[][] source, String[][] fix) {
		String[][] result = new String[4][4];
		for (int i = 0; i < 4; i++) {	// 控制行
			for (int j = 0; j < 4; j++) {	// 控制行使用次数
				String ret = "00000000";
				for (int k = 0; k < 4; k++) {	// 控制列
					String binMultiply = binMultiply(fix[i][k], source[k][j]);
					ret = xor(ret, binMultiply);
				}
				// 将结果转为16进制
				String hex = binToHex(ret);
				result[i][j] = hex;
			}
		}
		return result;
	}
	
	/**
	 * S盒替换
	 * @return
	 */
	public static String searchS(String hex, String[][] s) {
		/*
		 * 1. 将传入的2位16进制字符串前后位分开, 转成10进制
		 * 2. 第一位表示行, 第2位表示列, 去查询S盒, 返回2位16进制字符串
		 */
		int row = Integer.valueOf(hex.substring(0, 1), 16);
		int col = Integer.valueOf(hex.substring(1, 2), 16);
		return s[row][col];
	}
	
	/**
	 * 实现两个4 * 4矩阵对应元素的xor操作
	 * @return
	 */
	public static String[][] xor4HexArr(String[][] h1, String[][] h2) {
		String[][] ret = new String[4][4];
		for (int i = 0; i < 4; i++) {
			for (int j = 0; j < 4; j++) {
				ret[i][j] = xor4Hex(h1[i][j], h2[i][j]);
			}
		}
		return ret;
	}
	
	/**
	 * 将两个16进制字符串转成2进制字符串后进行乘法运算, 结果符合GF(2^8)
	 */
	public static String binMultiply(String hexStr1, String hexStr2) {
		// 1. 将16进制字串符转成二进制字符串, 不足8位则在前面补0
		String binStr1 = hexTobin(hexStr1);
		String binStr2 = hexTobin(hexStr2);
		
		// 2. 定义返回结果result, 遍历binStr1中字符, 遇到字符'1'时, 记录其索引index, 调用binMultiply(binStr2, 7 - index), 得到ret, 再与result进行xor运算
		String result = "00000000";
		for (int i = 0; i < binStr1.length(); i++) {
			if (binStr1.charAt(i) == '1') {
				String ret = binMultiply(binStr2, 7 - i);
				result = xor(result, ret);
			}
		}
		return result;
	}
	
	/**
	 * 计算x的n次方 * f(x)
	 * @param binStr f(x)二进制字符串
	 * @param num x的指数n
	 */
	public static String binMultiply(String binStr, int num) {
		/*
		 	思路:
		 		(1)循环执行num次以下的操作
		 		(2)是否flag判断最左边位是否为1
		 		(3)将binStr中二进制数左移一位, 右边补0
		 		(4)如果flag为true, 将(3)所得结果与00011011进行xor操作
		 */
		String ret = new String(binStr);
		for (int i = 0; i < num; i++) {
			boolean flag = ret.charAt(0) == '1';
			ret = ret.substring(1, ret.length()) + "0";
			if (flag) {
				ret = xor(ret, XOR_STR);
			}
		}
		return ret;
	}
	
	/**
	 * 两个十六进制字符串求异或
	 */
	public static String xor4Hex(String h1, String h2) {
		return binToHex(xor(hexTobin(h1), hexTobin(h2)));
	}
	
	/**
	 * 两个二进制字符串求异或
	 */
	public static String xor(String s1, String s2) {
		if (s1 == null || s2 == null || s1.length() != s2.length())
			throw new RuntimeException("字符串长度不相等");
		StringBuilder sb = new StringBuilder(s1.length());
		for (int i = 0; i < s1.length(); i++) {
			if (s1.charAt(i) == s2.charAt(i)) {
				sb.append("0");
			} else {
				sb.append("1");
			}
		}
		return sb.toString();
	}
	
	/**
	 * 8位二进制转2位十六进制
	 * @return
	 */
	public static String binToHex(String bin) {
		String hex = Integer.toHexString(Integer.valueOf(bin, 2));
		if (hex.length() == 1) {
			hex = "0" + hex;
		}
		return hex;
	}
	
	/**
	 * 2位十六进制转8位二进制
	 * @param hex
	 * @return
	 */
	public static String hexTobin(String hex) {
		String bin = Integer.toBinaryString(Integer.valueOf(hex, 16));
		int len1 = 8 - bin.length();
		if (len1 != 0) { 
			bin = ZEROS.substring(0, len1) + bin;
		}
		return bin;
	}
}
